import torch
from tqdm import tqdm
from utils import todevice


QUESTION_CATEGORY = {0:'count',1:'exist',2:'query_color',3:'query_size',4:'query_actiontype',5:'query_direction',
                     6:'query_shape',7:'compare_more',8:'compare_equal',9:'compare_less',10:'attribute_compare_color',
                     11:'attribute_compare_size',12:'attribute_compare_actiontype',13:'attribute_compare_direction',
                     14:'attribute_compare_shape'}


def validate(batch_size, model, data, device, write_preds=False):
    model.eval()
    print('validating...')
    total_acc, count = 0.0, 0
    all_preds = []
    gts = []
    v_ids = []
    q_ids = []
    what_acc,who_acc,how_acc,when_acc,where_acc = 0.,0.,0.,0.,0.
    what_count, who_count, how_count, when_count, where_count = 0,0,0,0,0

    with torch.no_grad():
        for batch in tqdm(data, total=len(data)): # tqdm是一个进度条工具，参数是可迭代对象
            video_ids, question_ids, answers, appearance_clips, motion_clips, question, question_len = [todevice(x, device) for x in batch]
            if batch_size == 1:
                answers = answers.to(device)
            else:
                answers = answers.to(device).squeeze()
            logits, aq_embed, mq_embed, com_app, com_motion, aq_fusion, mq_fusion = model(appearance_clips, motion_clips, question, question_len, stage=2)
            preds = logits.detach().argmax(1)
            agreeings = (preds == answers)
            what_idx = []
            who_idx = []
            how_idx = []
            when_idx = []
            where_idx = []

            key_word = question[:,0].to('cpu') # batch-based questions word
            for i,word in enumerate(key_word):
                word = int(word)
                if data.vocab['question_idx_to_token'][word] == 'what':
                    what_idx.append(i)
                elif data.vocab['question_idx_to_token'][word] == 'who':
                    who_idx.append(i)
                elif data.vocab['question_idx_to_token'][word] == 'how':
                    how_idx.append(i)
                elif data.vocab['question_idx_to_token'][word] == 'when':
                    when_idx.append(i)
                elif data.vocab['question_idx_to_token'][word] == 'where':
                    where_idx.append(i)

            if write_preds:
                preds = logits.argmax(1)
                answer_vocab = data.vocab['answer_idx_to_token']

                for predict in preds:
                    all_preds.append(answer_vocab[predict.item()])
                
                for gt in answers:
                    gts.append(answer_vocab[gt.item()])
                
                for id in video_ids:
                    v_ids.append(id.cpu().numpy())
                for ques_id in question_ids:
                    q_ids.append(ques_id)
            
            total_acc += agreeings.float().sum().item()
            count += answers.size(0)

            what_acc += agreeings.float()[what_idx].sum().item() if what_idx != [] else 0
            who_acc += agreeings.float()[who_idx].sum().item() if who_idx != [] else 0
            how_acc += agreeings.float()[how_idx].sum().item() if how_idx != [] else 0
            when_acc += agreeings.float()[when_idx].sum().item() if when_idx != [] else 0
            where_acc += agreeings.float()[where_idx].sum().item() if where_idx != [] else 0
            what_count += len(what_idx)
            who_count += len(who_idx)
            how_count += len(how_idx)
            when_count += len(when_idx)
            where_count += len(where_idx)

        acc = total_acc / count
        what_acc = what_acc / what_count
        who_acc = who_acc / who_count
        how_acc = how_acc / how_count
        when_acc = when_acc / when_count
        where_acc = where_acc / where_count

    if not write_preds:
        return acc, what_acc, who_acc, how_acc, when_acc, where_acc
    else:
        return acc, all_preds, gts, v_ids, q_ids, what_acc, who_acc, how_acc, when_acc, where_acc
    